from tools import *
from train import *


def eval_rnn(line_tensor):
    hidden = rnn.init_hidden()
    for i in range(line_tensor.size()[0]):
        output, hidden = rnn(line_tensor[i], hidden)
    return output.squeeze(0)


def eval_lstm(line_tensor):
    hidden, c = lstm.init_hidden_c()
    for i in range(line_tensor.size()[0]):
        output, hidden = rnn(line_tensor[i], hidden)
    return output.squeeze(0)


def eval_gru(line_tensor):
    hidden = gru.init_hidden()
    for i in range(line_tensor.size()[0]):
        output, hidden = rnn(line_tensor[i], hidden)
    return output.squeeze(0)


def eval_(input_line, eval_fn, n_predict=3):
    print(input_line)
    with torch.no_grad():
        output = eval_fn(line_to_tensor(input_line))

        topv, topi = output.topk(n_predict, 1, True)
        predict = []
        for i in range(n_predict):
            value = topv[0][i].item()
            category_index = topi[0][i].item()
            print('%2f %s' % (value, all_categories[category_index]))
            predict.append([value, all_categories[category_index]])


for eval_fn in [eval_rnn, eval_lstm, eval_gru]:
    print('-'*20)
    eval_('Wong', eval_fn)
    eval_('Jackson', eval_fn)
    eval_('Qi', eval_fn)
